import networkx as nx

from EmpVoting.Model import EmpNetwork
from EmpVoting.Model import Agent

class NetworkGen(object):


    def __init__(self,
                 numAgents,
                 numAlternatives,
                 preAgentGenList=None,
                 edgeGenList=None,
                 postAgentGenList=None):
        '''
        Constructor
        '''
        if preAgentGenList is None:
            preAgentGenList = []
        if edgeGenList is None:
            edgeGenList = []
        if postAgentGenList is None:
            postAgentGenList = []
        
        self.numAgents = numAgents
        self.numAlternatives = numAlternatives
        self.preAgentGenList = preAgentGenList
        self.edgeGenList = edgeGenList
        self.postAgentGenList = postAgentGenList
        
    def Generate(self):
        alternativeSet = set([])
        
        for i in range(0, self.numAlternatives):
            alternativeSet.add("a" + str(i))
        
        network = EmpNetwork(nx.DiGraph(), [], alternativeSet)
        
        for i in range(0, self.numAgents):
            network.agentList.append(Agent(i))
            
        network.graph.add_nodes_from(network.agentList)
        
        for preAgentGen in self.preAgentGenList:
            preAgentGen.Apply(network)
        
        for edgeGen in self.edgeGenList:
            edgeGen.Apply(network)
            
        for postAgentGen in self.postAgentGenList:
            postAgentGen.Apply(network)
            
        #Set the in degree and out degree for each agent.
        inDegreeDict = network.graph.in_degree()
        outDegreeDict = network.graph.out_degree()
        for agent in network.agentList:
            agent.inDegree = inDegreeDict[agent]
            agent.outDegree = outDegreeDict[agent]

        return network
    